#!/usr/bin/env python3
"""
修复脚本 - 删除有问题的索引，重新生成数据
"""

import psycopg2
import os
import requests
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

PROXY_HOST = "172.26.134.84"
PROXY_PORT = 7890
PROXIES = {
    'http': f'http://{PROXY_HOST}:{PROXY_PORT}',
    'https': f'http://{PROXY_HOST}:{PROXY_PORT}'
}

def get_db_connection():
    try:
        return psycopg2.connect(dbname="vectortutorialdb")
    except:
        try:
            return psycopg2.connect(dbname="vectortutorialdb", host="/var/run/postgresql")
        except:
            return psycopg2.connect(
                dbname="vectortutorialdb",
                user="postgres",
                password="postgres",
                host="localhost"
            )

class OpenAIEmbeddings:
    def __init__(self):
        with open("API_KEY.txt", "r") as f:
            self.api_key = f.read().strip()
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
    
    def embed_query(self, text):
        data = {
            "input": text,
            "model": "text-embedding-3-small"
        }
        
        response = requests.post(
            "https://api.openai.com/v1/embeddings",
            headers=self.headers,
            json=data,
            proxies=PROXIES,
            verify=False,
            timeout=30
        )
        
        if response.status_code == 200:
            result = response.json()
            return result['data'][0]['embedding']
        else:
            raise Exception(f"API错误: {response.status_code}")

texts = [
    "Apple MacBook Pro 14-inch with M3 chip, 16GB RAM, 1TB SSD, Space Gray",
    "Lenovo ThinkPad X1 Carbon, Intel i7, 32GB RAM, 1TB SSD, lightweight business laptop",
    "ASUS ROG Strix gaming laptop, AMD Ryzen 9, RTX 4070 GPU, 32GB RAM, RGB keyboard",
    "Dell XPS 13, Intel i5, 16GB RAM, 512GB SSD, ultra-slim design",
    "HP Envy desktop, AMD Ryzen 7, 64GB RAM, 2TB SSD, NVIDIA RTX 4060",
    "Samsung Galaxy Book3, Intel i7, 16GB RAM, AMOLED display, 1TB SSD",
    "Microsoft Surface Laptop 6, Intel i7, 16GB RAM, touchscreen, Platinum",
    "Acer Predator Orion desktop, Intel i9, RTX 4090 GPU, 128GB RAM, liquid cooling",
    "MSI Creator Z16, Intel i9, RTX 4080, 64GB RAM, designed for content creators",
    "Framework Laptop DIY Edition, modular components, Intel i7, 32GB RAM, 1TB SSD",
    "Apple iMac 27-inch, M3 chip, 16GB unified memory, 1TB SSD, 5K Retina display",
    "Raspberry Pi 5 Kit, 8GB RAM, 256GB SD card, micro PC for developers"
]

def fix_database():
    print("=" * 80)
    print("🔧 修复数据库")
    print("=" * 80)
    
    conn = get_db_connection()
    cur = conn.cursor()
    
    # 步骤1: 删除所有索引
    print("\n步骤 1: 删除有问题的索引...")
    try:
        # 查找所有索引
        cur.execute("""
            SELECT indexname 
            FROM pg_indexes 
            WHERE tablename = 'embeddings' 
            AND indexname != 'embeddings_pkey';
        """)
        indexes = cur.fetchall()
        
        for idx in indexes:
            idx_name = idx[0]
            print(f"   删除索引: {idx_name}")
            cur.execute(f"DROP INDEX IF EXISTS {idx_name};")
        
        conn.commit()
        print("   ✅ 索引已删除")
    except Exception as e:
        print(f"   警告: {e}")
        conn.rollback()
    
    # 步骤2: 清空表
    print("\n步骤 2: 清空表...")
    cur.execute("TRUNCATE TABLE embeddings;")
    conn.commit()
    print("   ✅ 表已清空")
    
    # 步骤3: 重新生成嵌入向量
    print("\n步骤 3: 重新生成嵌入向量...")
    embeddings_model = OpenAIEmbeddings()
    embeddings = []
    
    for i, text in enumerate(texts, 1):
        print(f"   [{i}/12] 生成: {text[:50]}...")
        vec = embeddings_model.embed_query(text)
        embeddings.append(vec)
    
    print("   ✅ 所有向量生成完成")
    
    # 步骤4: 插入数据（不创建索引）
    print("\n步骤 4: 插入数据...")
    for i, (vec, content) in enumerate(zip(embeddings, texts), 1):
        vec_str = '[' + ','.join(map(str, vec)) + ']'
        cur.execute(
            "INSERT INTO embeddings (content, embedding) VALUES (%s, %s::vector);",
            (content, vec_str)
        )
        if i % 3 == 0:
            print(f"   已插入 {i}/12 条...")
    
    conn.commit()
    print("   ✅ 数据插入完成")
    
    # 步骤5: 验证查询
    print("\n步骤 5: 验证查询...")
    
    # 测试查询1: Apple products
    test_query1 = "Apple products"
    print(f"\n   测试查询: {test_query1}")
    test_vec1 = embeddings_model.embed_query(test_query1)
    test_vec1_str = '[' + ','.join(map(str, test_vec1)) + ']'
    
    cur.execute("""
        SELECT id, content, embedding <=> %s::vector AS distance
        FROM embeddings
        ORDER BY distance
        LIMIT 5;
    """, (test_vec1_str,))
    
    results1 = cur.fetchall()
    print(f"   返回 {len(results1)} 条结果:")
    for r in results1:
        print(f"      ID{r[0]}: {r[1][:50]}... (距离: {r[2]:.4f})")
    
    # 测试查询2: gaming laptop
    test_query2 = "gaming laptop with RGB keyboard"
    print(f"\n   测试查询: {test_query2}")
    test_vec2 = embeddings_model.embed_query(test_query2)
    test_vec2_str = '[' + ','.join(map(str, test_vec2)) + ']'
    
    cur.execute("""
        SELECT id, content, embedding <=> %s::vector AS distance
        FROM embeddings
        ORDER BY distance
        LIMIT 5;
    """, (test_vec2_str,))
    
    results2 = cur.fetchall()
    print(f"   返回 {len(results2)} 条结果:")
    for r in results2:
        print(f"      ID{r[0]}: {r[1][:50]}... (距离: {r[2]:.4f})")
    
    cur.close()
    conn.close()
    
    print("\n" + "=" * 80)
    print("✅ 修复完成！")
    print("=" * 80)
    print("\n注意：")
    print("  - 已删除有问题的索引")
    print("  - 数据量较少时不需要索引")
    print("  - 现在可以正常使用了")
    print("\n运行主程序测试:")
    print("  python test_pgvector_wsl.py")

if __name__ == "__main__":
    fix_database()